package com.example.service.impl;

import com.example.config.AiConfig;
import com.example.model.Message;
import com.example.service.AiService;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import okhttp3.*;
import okio.BufferedSource;
import org.springframework.stereotype.Service;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import com.fasterxml.jackson.databind.JsonNode;

@Slf4j
@Service
@RequiredArgsConstructor
public class GuijiAiService implements AiService {
    private final AiConfig aiConfig;
    private final ObjectMapper objectMapper;
    private final OkHttpClient client = new OkHttpClient.Builder()
        .connectTimeout(60, TimeUnit.SECONDS)
        .readTimeout(60, TimeUnit.SECONDS)
        .writeTimeout(60, TimeUnit.SECONDS)
        .retryOnConnectionFailure(true)
        .build();
    private final ExecutorService executorService = Executors.newCachedThreadPool();

    // 使用Map来存储所有会话的消息历史
    private final Map<String, List<Message>> messageHistory = new HashMap<>();

    @Override
    public String generateResponse(String prompt, String sessionId) {
        try {
            // 如果sessionId为空，创建新的会话
            final String finalSessionId = (sessionId == null || sessionId.isEmpty()) ? 
                UUID.randomUUID().toString() : sessionId;
            
            // 获取或创建当前会话的消息历史
            List<Message> history = messageHistory.computeIfAbsent(finalSessionId, k -> new ArrayList<>());
            
            // 如果是新会话，添加系统角色设定
            if (history.isEmpty()) {
                history.add(new Message("system", "你是一位经验丰富的大学老师名字叫小软，擅长用引导式教学方式帮助学生理解问题。" +
                    "你的回答应该：\n" +
                    "1. 用温和、鼓励的语气\n" +
                    "2. 通过提问引导学生思考\n" +
                    "3. 用生动的例子解释复杂概念\n" +
                    "4. 适当使用比喻和类比\n" +
                    "5. 在回答结束时，提出一个相关的问题让学生继续思考\n" +
                    "6. 保持专业性和学术严谨性"));
            }
            
            // 添加用户的新消息
            history.add(new Message("user", prompt));

            // 构建完整的消息列表
            List<Map<String, String>> messages = new ArrayList<>();
            for (Message message : history) {
                Map<String, String> messageMap = new HashMap<>();
                messageMap.put("role", message.getRole());
                messageMap.put("content", message.getContent());
                messages.add(messageMap);
            }

            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", "Qwen/QwQ-32B");
            requestBody.put("messages", messages);

            String jsonBody = objectMapper.writeValueAsString(requestBody);
            log.info("发送请求体: {}", jsonBody);
            
            RequestBody body = RequestBody.create(
                MediaType.parse("application/json"),
                jsonBody
            );

            Request request = new Request.Builder()
                .url(aiConfig.getGuiji().getApiUrl())
                .addHeader("Authorization", "Bearer " + aiConfig.getGuiji().getApiKey())
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();

            try (Response response = client.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    String errorBody = response.body() != null ? response.body().string() : "无错误详情";
                    log.error("API调用失败: HTTP {} - {}", response.code(), errorBody);
                    throw new RuntimeException("API调用失败: " + response.code() + " - " + errorBody);
                }

                String responseBody = response.body().string();
                Map<String, Object> responseMap = objectMapper.readValue(responseBody, Map.class);
                
                @SuppressWarnings("unchecked")
                List<Map<String, Object>> choices = (List<Map<String, Object>>) responseMap.get("choices");
                if (choices == null || choices.isEmpty()) {
                    throw new RuntimeException("API响应中没有choices数据");
                }
                
                Map<String, Object> firstChoice = choices.get(0);
                Map<String, Object> message = (Map<String, Object>) firstChoice.get("message");
                if (message == null) {
                    throw new RuntimeException("API响应中没有message数据");
                }
                
                String content = (String) message.get("content");
                if (content == null) {
                    throw new RuntimeException("API响应中没有content数据");
                }
                
                // 添加AI的回复到历史记录
                history.add(new Message("assistant", content));
                
                return content;
            }
        } catch (Exception e) {
            log.error("生成响应时发生错误", e);
            throw new RuntimeException("生成响应时发生错误: " + e.getMessage(), e);
        }
    }

    @Override
    public Flux<String> generateStreamResponse(String prompt, String sessionId) {
        return Flux.create(sink -> {
            executorService.execute(() -> {
                try {
                    // 如果sessionId为空，创建新的会话
                    final String finalSessionId = (sessionId == null || sessionId.isEmpty()) ? 
                        UUID.randomUUID().toString() : sessionId;
                    
                    // 获取或创建当前会话的消息历史
                    List<Message> history = messageHistory.computeIfAbsent(finalSessionId, k -> new ArrayList<>());
                    
                    // 如果是新会话，添加系统角色设定
                    if (history.isEmpty()) {
                        history.add(new Message("system", "你是一位经验丰富的大学老师名字叫小软，擅长用引导式教学方式帮助学生理解问题。" +
                            "你的回答应该：\n" +
                            "1. 用温和、鼓励的语气\n" +
                            "2. 通过提问引导学生思考\n" +
                            "3. 用生动的例子解释复杂概念\n" +
                            "4. 适当使用比喻和类比\n" +
                            "5. 在回答结束时，提出一个相关的问题让学生继续思考\n" +
                            "6. 保持专业性和学术严谨性"));
                    }
                    
                    // 添加用户的新消息
                    history.add(new Message("user", prompt));

                    // 构建完整的消息列表
                    List<Map<String, String>> messages = new ArrayList<>();
                    for (Message message : history) {
                        Map<String, String> messageMap = new HashMap<>();
                        messageMap.put("role", message.getRole());
                        messageMap.put("content", message.getContent());
                        messages.add(messageMap);
                    }

                    Map<String, Object> requestBody = new HashMap<>();
                    requestBody.put("model", "Qwen/QwQ-32B");
                    requestBody.put("messages", messages);
                    requestBody.put("stream", true);

                    String jsonBody = objectMapper.writeValueAsString(requestBody);
                    log.info("发送流式请求体: {}", jsonBody);
                    
                    RequestBody body = RequestBody.create(
                        MediaType.parse("application/json"),
                        jsonBody
                    );

                    Request request = new Request.Builder()
                        .url(aiConfig.getGuiji().getApiUrl())
                        .addHeader("Authorization", "Bearer " + aiConfig.getGuiji().getApiKey())
                        .addHeader("Content-Type", "application/json")
                        .post(body)
                        .build();

                    StringBuilder fullResponse = new StringBuilder();
                    try (Response response = client.newCall(request).execute()) {
                        if (!response.isSuccessful()) {
                            String errorBody = response.body() != null ? response.body().string() : "无错误详情";
                            log.error("API调用失败: HTTP {} - {}", response.code(), errorBody);
                            Map<String, Object> errorResponse = new HashMap<>();
                            errorResponse.put("code", response.code());
                            errorResponse.put("error", "API调用失败: " + response.code() + " - " + errorBody);
                            sink.next(objectMapper.writeValueAsString(errorResponse));
                            return;
                        }

                        ResponseBody responseBody = response.body();
                        if (responseBody == null) {
                            Map<String, Object> errorResponse = new HashMap<>();
                            errorResponse.put("code", 500);
                            errorResponse.put("error", "响应体为空");
                            sink.next(objectMapper.writeValueAsString(errorResponse));
                            return;
                        }

                        try (BufferedSource source = responseBody.source()) {
                            while (!source.exhausted()) {
                                String line = source.readUtf8Line();
                                if (line == null) continue;
                                
                                if (line.startsWith("data: ")) {
                                    String data = line.substring(6);
                                    
                                    if (data.equals("[DONE]")) {
                                        // 添加AI的完整回复到历史记录
                                        String finalResponse = fullResponse.toString();
                                        history.add(new Message("assistant", finalResponse));
                                        Map<String, Object> doneResponse = new HashMap<>();
                                        doneResponse.put("code", 200);
                                        doneResponse.put("done", true);
                                        sink.next(objectMapper.writeValueAsString(doneResponse));
                                        break;
                                    }
                                    
                                    try {
                                        Map<String, Object> responseMap = objectMapper.readValue(data, Map.class);
                                        @SuppressWarnings("unchecked")
                                        List<Map<String, Object>> choices = (List<Map<String, Object>>) responseMap.get("choices");
                                        if (choices != null && !choices.isEmpty()) {
                                            Map<String, Object> firstChoice = choices.get(0);
                                            Map<String, Object> delta = (Map<String, Object>) firstChoice.get("delta");
                                            if (delta != null) {
                                                String content = (String) delta.get("content");
                                                if (content != null) {
                                                    fullResponse.append(content);
                                                    Map<String, Object> responseData = new HashMap<>();
                                                    responseData.put("code", 200);
                                                    responseData.put("content", content);
                                                    sink.next(objectMapper.writeValueAsString(responseData));
                                                }
                                            }
                                        }
                                    } catch (Exception e) {
                                        log.error("解析流式响应数据失败: {}", e.getMessage());
                                        Map<String, Object> errorResponse = new HashMap<>();
                                        errorResponse.put("code", 500);
                                        errorResponse.put("error", "解析流式响应数据失败: " + e.getMessage());
                                        sink.next(objectMapper.writeValueAsString(errorResponse));
                                    }
                                }
                            }
                        }
                    }
                } catch (Exception e) {
                    log.error("流式响应处理失败", e);
                    Map<String, Object> errorResponse = new HashMap<>();
                    errorResponse.put("code", 500);
                    errorResponse.put("error", "流式响应处理失败: " + e.getMessage());
                    sink.next(objectMapper.writeValueAsString(errorResponse));
                } finally {
                    sink.complete();
                }
            });
        });
    }

    // 清理会话历史的方法
    public void clearHistory(String sessionId) {
        messageHistory.remove(sessionId);
    }
} 